package wikiextract.nlp.trainingset.x;

import static wikiextract.nlp.trainingset.x.Settings.CROSS_VALIDATION_SETS;

import java.io.IOException;

import learning.crf.model.Model;
import learning.crf.training.BasicParseScorer;
import learning.crf.training.CRF;
import learning.crf.training.CRFParameters;
import learning.crf.training.CRF.DecodingResult;
import learning.data.Dataset;
import learning.data.SequenceDataReader;
import learning.data.document.SequenceDocument;

public class TrainCRF {

	static String tmpDir = "/projects/pardosa/s2/raphaelh/tmp";
	static int REPETITIONS = 1;
	
	public static double[] getCrossValidationResults() 
		throws IOException {
		
		Eval eval = new Eval();
		for (int i=0; i < CROSS_VALIDATION_SETS; i++) {
			String s = "_" + i;
			// load datasets
			Dataset<SequenceDocument> trainData = 
				SequenceDataReader.read(tmpDir + "/train" + s, "utf-8");
			Dataset<SequenceDocument> testData = 
				SequenceDataReader.read(tmpDir + "/test" + s, "utf-8"); 
			//System.out.println(trainData.numDocs() + " train, " + testData.numDocs() + " test");
			
			for (int r=0; r < REPETITIONS; r++) {
				// running cross validation			
				learning.crf.training.CollinsTraining ct = new learning.crf.training.CollinsTraining();
				CRFParameters params = ct.train(trainData);
				test(testData, params, eval);
			}
			//if (r==0)
			//	generateResults(dir + "/results_new" + s + suffix, evaluatorCV);
		}
		double precision = eval.avgPrecision();
		double recall = eval.avgRecall();
		double f1 = 2.0* (precision*recall) / (precision + recall);
		if (Double.isNaN(f1)) f1 = 0;
		
		//System.out.println("   CV AVG f1 " + f1 + ", precision " + eval.avgPrecision() + ", recall " + eval.avgRecall());
		return new double[] { f1, precision, recall };
	}

	public static void test(Dataset<SequenceDocument> testData, CRFParameters parameters, Eval eval) {
		parameters.model = new Model(testData.numLabels());

		BasicParseScorer allParseScorer = new BasicParseScorer(BasicParseScorer.Type.ALL_TERMINALS);
		allParseScorer.setDataset(testData);
		
		int truePos = 0; // articles (or words?)
		int falsePos = 0;
		int trueNeg = 0;
		int falseNeg = 0;
		int incorrectLabel = 0;
		
		String extracted = null;
		String labeled = null;
		int prevArticleId = -1;
		int articleId = -1;
		for (SequenceDocument doc : testData) {
			articleId = Integer.parseInt(doc.meta.get("articleId"));

			if (articleId != prevArticleId && prevArticleId != -1) {
				if (extracted == null && labeled != null) falseNeg++;
				if (extracted != null && labeled == null) falsePos++;
				if (extracted == null && labeled == null) trueNeg++;
				if (extracted != null && labeled != null &&
						extracted.equals(labeled)) truePos++;
				if (extracted != null && labeled != null &&
						!extracted.equals(labeled)) incorrectLabel++;
				/*				
				if (extracted != null || labeled != null) {
					System.out.println(" predicted: " + extracted);
					System.out.println(" actual:    " + labeled);
				}*/
				extracted = null;
				labeled = null;
			}
				
			// compute most likely label under current parameters
			DecodingResult predictedParse = CRF.decode(doc, allParseScorer, parameters);
			DecodingResult trueParse = new DecodingResult(doc.labels, 0);
			
			if (extracted == null)
				for (int i=0; i < predictedParse.labels.length; i++) {
					if (predictedParse.labels[i] == 1) {
						if (extracted == null) extracted = "";
						extracted += doc.tokens[i] + " ";
					}
				}
			
			if (labeled == null)
				for (int i=0; i < trueParse.labels.length; i++) {
					if (trueParse.labels[i] == 1) {
						if (labeled == null) labeled = "";
						labeled += doc.tokens[i] + " ";
					}
				}
			
			
			// update evaluator
			//evaluator.update(doc, predictedParse, trueParse);
			prevArticleId = articleId;
		}
		// last one
		if (prevArticleId != -1) {
			if (extracted == null && labeled != null) falseNeg++;
			if (extracted != null && labeled == null) falsePos++;
			if (extracted == null && labeled == null) trueNeg++;
			if (extracted != null && labeled != null &&
					extracted.equals(labeled)) truePos++;
			if (extracted != null && labeled != null &&
					!extracted.equals(labeled)) incorrectLabel++;
			//System.out.println("labeled " + labeled);
		}
		
		double precision = truePos / (double)(incorrectLabel + truePos + falsePos);
		double recall = truePos / (double)(truePos + falseNeg + incorrectLabel);
		if (testData.numDocs() == 0) System.out.println("warning: test data size = 0");
		if (Double.isNaN(precision)) precision = 1; // not
		if (Double.isNaN(recall)) recall = 1; // not
		eval.update(precision, recall);
		//System.out.println("    precision " + precision + ", recall " + recall);
	}
	
	static class Eval {
		int num = 0;
		double precision = 0;
		double recall = 0;
		
		public void update(double precision, double recall) {
			this.precision += precision;
			this.recall += recall;
			this.num++;
		}
		
		public double avgPrecision() {
			return precision / num;
		}
		
		public double avgRecall() {
			return recall / num;
		}
	}
}
